import argparse
from LSTMGanForShell import LSTMGanForShell as GanForShell
import yaml
import os

path = 'src/config.yaml'

if __name__ == '__main__':
    with open(os.path.join(os.getcwd(),path), 'r') as config_yaml:
        data = yaml.load(config_yaml, yaml.FullLoader)

    parser = argparse.ArgumentParser()
    parser.add_argument('--cuda', default=data['cuda'], type=int)
    parser.add_argument('--start', default=data['start'], type=int)
    parser.add_argument('--load_model', default=data['load_model'], type=int)

    parser.add_argument('--embed_dim', default=data['embed_dim'], type=int)
    parser.add_argument('--ff1_out', default=data['ff1_out'], type=int)
    parser.add_argument('--hidden_dim', default=data['hidden_dim'], type=int)
    parser.add_argument('--generator_out_dim', default=data['generator_out_dim'], type=int)
    parser.add_argument('--discriminator_out_dim', default=data['discriminator_out_dim'], type=int)

    parser.add_argument('--lyrics_dis_rate', default=data['lyrics_dis_rate'], type=float)
    parser.add_argument('--learning_rate_G', default=data['learning_rate_G'], type=float)
    parser.add_argument('--learning_rate_D', default=data['learning_rate_D'], type=float)

    parser.add_argument('--start_epoch', default=data['start_epoch'], type=int)

    parser.add_argument('--epochs', default=data['epochs'], type=int)
    parser.add_argument('--save_per_epoch', default=data['save_per_epoch'], type=int)
    parser.add_argument('--train_D_steps', default=data['train_D_steps'], type=int)
    parser.add_argument('--train_G_steps', default=data['train_G_steps'], type=int)
    parser.add_argument('--sequence_len', default=data['sequence_len'], type=int)

    parser.add_argument('--mem_slots', default=data['mem_slots'], type=int)
    parser.add_argument('--num_heads', default=data['num_heads'], type=int)
    parser.add_argument('--head_size', default=data['head_size'], type=int)

    parser.add_argument('--EPS', default=data['EPS'], type=float)

    parser.add_argument('--batch_size', default=data['data_params']['batch_size'], type=int)
    parser.add_argument('--shuffle', default=data['data_params']['shuffle'], type=int)
    parser.add_argument('--num_workers', default=data['data_params']['num_workers'], type=int)

    opt = parser.parse_args()
    # print(list(opt))
    model = GanForShell(opt)
    model.train_all_epochs()
